//
// Created on 2025/6/15.
//
// Node APIs are not fully supported. To solve the compilation error of the interface cannot be found,
// please include "napi/native_api.h".

#ifndef VARINT_COMMON_H
#define VARINT_COMMON_H
#include "stdint.h"
#include <cstddef>

#define MSB 0x80
#define REST 0x7f

enum class OP_RESULT { SUCCESS, BAD_DATA, OUTPUT_TOO_SMALL, NEED_MORE_INPUT };

static const uint8_t len_tab[] = {0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3,
                                  4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 7,
                                  7, 7, 7, 7, 7, 7, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 10};

#ifdef __aarch64__
#include <arm_neon.h>
inline uint32x4_t varint_bytes_vec(uint32x4_t value) {
    value = vorrq_u32(value, vdupq_n_u32(1));
    uint32x4_t lz = vclzq_u32(value);
    const uint32x4_t max_bits = vdupq_n_u32(32);
    uint32x4_t msb_pos = vsubq_u32(max_bits, lz);

    uint8x16x2_t table = vld1q_u8_x2(len_tab);
    uint8x16_t indices = vreinterpretq_u8_u32(msb_pos);

    uint8x16_t result_u8 = vqtbl2q_u8(table, indices);

    return vreinterpretq_u32_u8(result_u8);
}

inline uint32x4_t vzigzag_s32(int32x4_t value) {
    int32x4_t shifted = vshlq_n_s32(value, 1);
    int32x4_t sign = vshrq_n_s32(value, 31);
    return veorq_u32(vreinterpretq_u32_s32(shifted), sign);
}

inline uint32x4_t varint_bytes_vec(int32x4_t value) {
    uint32x4_t val = vzigzag_s32(value);
    return varint_bytes_vec(val);
}

inline uint8x16_t varint_count_vec(uint8x16_t value) { return vshrq_n_u8(vcltq_u8(value, vdupq_n_u8(0x80)), 7); }


#endif


inline uint64_t zigzag_encode(int64_t n) { return ((uint64_t)n << 1) ^ ((int64_t)(n >> 63)); }
inline uint32_t zigzag_encode(int32_t n) { return ((uint32_t)n << 1) ^ ((int32_t)(n >> 31)); }
inline uint64_t zigzag_encode(uint64_t n) { return n; }
inline uint32_t zigzag_encode(uint32_t n) { return n; }

inline int32_t zigzag_decode(uint32_t n) { return (n >> 1) ^ -(int32_t)(n & 1); }
inline int64_t zigzag_decode(uint64_t n) { return (n >> 1) ^ -(int64_t)(n & 1); }

inline uint32_t varint_bytes(uint64_t value) {
    if (value < MSB)
        return 1;

    return len_tab[64 - __builtin_clzll(value)];
}

inline uint32_t varint_bytes(uint32_t value) {
    if (value < MSB)
        return 1;

    return len_tab[32 - __builtin_clz((uint64_t)value)];
}

inline uint32_t varint_bytes(int64_t value) { return varint_bytes(zigzag_encode(value)); }
inline uint32_t varint_bytes(int32_t value) { return varint_bytes(zigzag_encode(value)); }


/**
 * 计算u32 varint编码后长度
 * @param input
 * @param length
 * @return
 */
static size_t varint_bytes_total(const uint32_t *input, size_t length) {
    size_t i = 0;
    size_t total = 0;
#ifdef __aarch64__
    // arm_v8a 使用neon加速
    for (; i + 32 <= length; i += 32) {
        uint32x4x4_t vecs0 = vld1q_u32_x4(input + i);
        uint32x4x4_t vecs1 = vld1q_u32_x4(input + i + 16);

        // 并行计算所有向量的字节数
        uint32x4_t bytes0 = varint_bytes_vec(vecs0.val[0]);
        uint32x4_t bytes1 = varint_bytes_vec(vecs0.val[1]);
        uint32x4_t bytes2 = varint_bytes_vec(vecs0.val[2]);
        uint32x4_t bytes3 = varint_bytes_vec(vecs0.val[3]);
        uint32x4_t bytes4 = varint_bytes_vec(vecs1.val[0]);
        uint32x4_t bytes5 = varint_bytes_vec(vecs1.val[1]);
        uint32x4_t bytes6 = varint_bytes_vec(vecs1.val[2]);
        uint32x4_t bytes7 = varint_bytes_vec(vecs1.val[3]);

        // 分层累加减少依赖
        uint32x4_t sum01 = vaddq_u32(bytes0, bytes1);
        uint32x4_t sum23 = vaddq_u32(bytes2, bytes3);
        uint32x4_t sum45 = vaddq_u32(bytes4, bytes5);
        uint32x4_t sum67 = vaddq_u32(bytes6, bytes7);

        uint32x4_t sum0123 = vaddq_u32(sum01, sum23);
        uint32x4_t sum4567 = vaddq_u32(sum45, sum67);

        uint32x4_t final_sum = vaddq_u32(sum0123, sum4567);

        total += vaddvq_u32(final_sum);
    }
    for (; i + 16 <= length; i += 16) {
        uint32x4x4_t vecs = vld1q_u32_x4(input + i);

        uint32x4_t bytes0 = varint_bytes_vec(vecs.val[0]);
        uint32x4_t bytes1 = varint_bytes_vec(vecs.val[1]);
        uint32x4_t bytes2 = varint_bytes_vec(vecs.val[2]);
        uint32x4_t bytes3 = varint_bytes_vec(vecs.val[3]);

        uint32x4_t sum_vec = vaddq_u32(vaddq_u32(bytes0, bytes1), vaddq_u32(bytes2, bytes3));
        total += vaddvq_u32(sum_vec);
    }
    for (; i + 4 <= length; i += 4) {
        uint32x4_t vec = vld1q_u32(input + i);
        uint32x4_t bytes_vec = varint_bytes_vec(vec);
        total += vaddvq_u32(bytes_vec);
    }

#endif

    // 降级默认处理
    for (; i < length; i++) {
        total += varint_bytes(input[i]);
    }
    return total;
}

/**
 * 计算s32 zigzag/varint编码后长度
 * @param input
 * @param length
 * @return
 */
static size_t varint_bytes_total(const int32_t *input, size_t length) {
    size_t i = 0;
    size_t total = 0;
#ifdef __aarch64__
    // arm_v8a 使用neon加速
    for (; i + 32 <= length; i += 32) {
        int32x4x4_t vecs0 = vld1q_s32_x4(input + i);
        int32x4x4_t vecs1 = vld1q_s32_x4(input + i + 16);

        // 并行计算所有向量的字节数
        uint32x4_t bytes0 = varint_bytes_vec(vecs0.val[0]);
        uint32x4_t bytes1 = varint_bytes_vec(vecs0.val[1]);
        uint32x4_t bytes2 = varint_bytes_vec(vecs0.val[2]);
        uint32x4_t bytes3 = varint_bytes_vec(vecs0.val[3]);
        uint32x4_t bytes4 = varint_bytes_vec(vecs1.val[0]);
        uint32x4_t bytes5 = varint_bytes_vec(vecs1.val[1]);
        uint32x4_t bytes6 = varint_bytes_vec(vecs1.val[2]);
        uint32x4_t bytes7 = varint_bytes_vec(vecs1.val[3]);

        // 分层累加减少依赖
        uint32x4_t sum01 = vaddq_u32(bytes0, bytes1);
        uint32x4_t sum23 = vaddq_u32(bytes2, bytes3);
        uint32x4_t sum45 = vaddq_u32(bytes4, bytes5);
        uint32x4_t sum67 = vaddq_u32(bytes6, bytes7);

        uint32x4_t sum0123 = vaddq_u32(sum01, sum23);
        uint32x4_t sum4567 = vaddq_u32(sum45, sum67);

        uint32x4_t final_sum = vaddq_u32(sum0123, sum4567);

        total += vaddvq_u32(final_sum);
    }
    for (; i + 16 <= length; i += 16) {
        int32x4x4_t vecs = vld1q_s32_x4(input + i);

        uint32x4_t bytes0 = varint_bytes_vec(vecs.val[0]);
        uint32x4_t bytes1 = varint_bytes_vec(vecs.val[1]);
        uint32x4_t bytes2 = varint_bytes_vec(vecs.val[2]);
        uint32x4_t bytes3 = varint_bytes_vec(vecs.val[3]);

        uint32x4_t sum_vec = vaddq_u32(vaddq_u32(bytes0, bytes1), vaddq_u32(bytes2, bytes3));
        total += vaddvq_u32(sum_vec);
    }
    for (; i + 4 <= length; i += 4) {
        int32x4_t vec = vld1q_s32(input + i);
        uint32x4_t bytes_vec = varint_bytes_vec(vec);
        total += vaddvq_u32(bytes_vec);
    }

#endif

    // 降级默认处理
    for (; i < length; i++) {
        total += varint_bytes(input[i]);
    }
    return total;
}


/**
 * 计算u64 varint编码长度
 * @param input
 * @param length
 * @return
 */
static size_t varint_bytes_total(const uint64_t *input, size_t length) {
    size_t total = 0;
#define GET_BYTES(index) total += varint_bytes(input[index])
    const uint64_t *end = input + length;

    while (input + 16 <= end) {
        GET_BYTES(0);
        GET_BYTES(1);
        GET_BYTES(2);
        GET_BYTES(3);
        GET_BYTES(4);
        GET_BYTES(5);
        GET_BYTES(6);
        GET_BYTES(7);
        GET_BYTES(8);
        GET_BYTES(9);
        GET_BYTES(10);
        GET_BYTES(11);
        GET_BYTES(12);
        GET_BYTES(13);
        GET_BYTES(14);
        GET_BYTES(15);
        input += 16;
    }

    while (input + 8 <= end) {
        GET_BYTES(0);
        GET_BYTES(1);
        GET_BYTES(2);
        GET_BYTES(3);
        GET_BYTES(4);
        GET_BYTES(5);
        GET_BYTES(6);
        GET_BYTES(7);
        input += 8;
    }

#undef GET_BYTES
    while (input < end) {
        total += varint_bytes(*input++);
    }

    return total;
}

/**
 * 计算u64 varint编码长度
 * @param input
 * @param length
 * @return
 */
static size_t varint_bytes_total(const int64_t *input, size_t length) {
    size_t total = 0;
#define GET_BYTES(index) total += varint_bytes(input[index])

    const int64_t *end = input + length;

    while (input + 16 <= end) {
        GET_BYTES(0);
        GET_BYTES(1);
        GET_BYTES(2);
        GET_BYTES(3);
        GET_BYTES(4);
        GET_BYTES(5);
        GET_BYTES(6);
        GET_BYTES(7);
        GET_BYTES(8);
        GET_BYTES(9);
        GET_BYTES(10);
        GET_BYTES(11);
        GET_BYTES(12);
        GET_BYTES(13);
        GET_BYTES(14);
        GET_BYTES(15);
        input += 16;
    }

    while (input + 8 <= end) {
        GET_BYTES(0);
        GET_BYTES(1);
        GET_BYTES(2);
        GET_BYTES(3);
        GET_BYTES(4);
        GET_BYTES(5);
        GET_BYTES(6);
        GET_BYTES(7);
        input += 8;
    }

#undef GET_BYTES
    while (input < end) {
        total += varint_bytes(*input++);
    }

    return total;
}

static size_t varint_count(const uint8_t *input, size_t length) {
    size_t i = 0;
    size_t total = 0;
#ifdef __aarch64__
    // arm_v8a 使用neon加速
    for (; i + 128 <= length; i += 128) {
        uint8x16x4_t vecs0 = vld1q_u8_x4(input + i);
        uint8x16x4_t vecs1 = vld1q_u8_x4(input + i + 64);

        // 并行计算所有向量的字节数
        uint8x16_t bytes0 = varint_count_vec(vecs0.val[0]);
        uint8x16_t bytes1 = varint_count_vec(vecs0.val[1]);
        uint8x16_t bytes2 = varint_count_vec(vecs0.val[2]);
        uint8x16_t bytes3 = varint_count_vec(vecs0.val[3]);
        uint8x16_t bytes4 = varint_count_vec(vecs1.val[0]);
        uint8x16_t bytes5 = varint_count_vec(vecs1.val[1]);
        uint8x16_t bytes6 = varint_count_vec(vecs1.val[2]);
        uint8x16_t bytes7 = varint_count_vec(vecs1.val[3]);

        // 分层累加减少依赖
        uint8x16_t sum01 = vaddq_u8(bytes0, bytes1);
        uint8x16_t sum23 = vaddq_u8(bytes2, bytes3);
        uint8x16_t sum45 = vaddq_u8(bytes4, bytes5);
        uint8x16_t sum67 = vaddq_u8(bytes6, bytes7);

        uint8x16_t sum0123 = vaddq_u8(sum01, sum23);
        uint8x16_t sum4567 = vaddq_u8(sum45, sum67);

        uint8x16_t final_sum = vaddq_u8(sum0123, sum4567);

        total += vaddvq_u8(final_sum);
    }
    for (; i + 64 <= length; i += 64) {
        uint8x16x4_t vecs = vld1q_u8_x4(input + i);

        uint8x16_t bytes0 = varint_count_vec(vecs.val[0]);
        uint8x16_t bytes1 = varint_count_vec(vecs.val[1]);
        uint8x16_t bytes2 = varint_count_vec(vecs.val[2]);
        uint8x16_t bytes3 = varint_count_vec(vecs.val[3]);

        uint8x16_t sum_vec = vaddq_u8(vaddq_u8(bytes0, bytes1), vaddq_u8(bytes2, bytes3));
        total += vaddvq_u8(sum_vec);
    }
    for (; i + 32 <= length; i += 32) {
        uint8x16x2_t vecs = vld1q_u8_x2(input + i);

        uint8x16_t bytes0 = varint_count_vec(vecs.val[0]);
        uint8x16_t bytes1 = varint_count_vec(vecs.val[1]);

        total += vaddvq_u8(vaddq_u8(bytes0, bytes1));
    }
    for (; i + 16 <= length; i += 16) {
        uint8x16_t vecs = vld1q_u8(input);
        total += vaddvq_u8(varint_count_vec(vecs));
    }

#endif

    // 降级默认处理
    for (; i < length; i++) {
        total += (input[i] & 0x80) ? 0 : 1;
    }
    return total;
}


#endif // VARINT_COMMON_H
